{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "6d1ca9ac",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/llm/bedrock_converse.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e3a8796-edc8-43f2-94ad-fe4fb20d70ed",
   "metadata": {},
   "source": [
    "# Bedrock Converse"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b007403c-6b7a-420c-92f1-4171d05ed9bb",
   "metadata": {},
   "source": [
    "## Basic Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ead155e-b8bd-46f9-ab9b-28fc009361dd",
   "metadata": {},
   "source": [
    "#### Call `complete` with a prompt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "211db735",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f250bb49",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-bedrock-converse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81732d83",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60be18ae-c957-4ac2-a58a-0652e18ee6d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.bedrock_converse import BedrockConverse\n",
    "\n",
    "profile_name = \"Your aws profile name\"\n",
    "resp = BedrockConverse(\n",
    "    model=\"anthropic.claude-3-haiku-20240307-v1:0\",\n",
    "    profile_name=profile_name,\n",
    ").complete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac2cbebb-a444-4a46-9d85-b265a3483d68",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14831268-f90f-499d-9d86-925dbc88292b",
   "metadata": {},
   "source": [
    "#### Call `chat` with a list of messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbe29574-4af1-48d5-9739-f60652b6ce6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.llms import ChatMessage\n",
    "from llama_index.llms.bedrock_converse import BedrockConverse\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"Tell me a story\"),\n",
    "]\n",
    "\n",
    "resp = BedrockConverse(\n",
    "    model=\"anthropic.claude-3-haiku-20240307-v1:0\",\n",
    "    profile_name=profile_name,\n",
    ").chat(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cbd550a-0264-4a11-9b2c-a08d8723a5ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ed5e894-4597-4911-a623-591560f72b82",
   "metadata": {},
   "source": [
    "## Streaming"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cb7986f-aaed-42e2-abdd-f274f6d4fc59",
   "metadata": {},
   "source": [
    "Using `stream_complete` endpoint "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d43f17a2-0aeb-464b-a7a7-732ba5e8ef24",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.bedrock_converse import BedrockConverse\n",
    "\n",
    "llm = BedrockConverse(\n",
    "    model=\"anthropic.claude-3-haiku-20240307-v1:0\",\n",
    "    profile_name=profile_name,\n",
    ")\n",
    "resp = llm.stream_complete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0214e911-cf0d-489c-bc48-9bb1d8bf65d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "for r in resp:\n",
    "    print(r.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40350dd8-3f50-4a2f-8545-5723942039bb",
   "metadata": {},
   "source": [
    "Using `stream_chat` endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc636e65-a67b-4dcd-ac60-b25abc9d8dbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.bedrock_converse import BedrockConverse\n",
    "\n",
    "llm = BedrockConverse(\n",
    "    model=\"anthropic.claude-3-haiku-20240307-v1:0\",\n",
    "    profile_name=profile_name,\n",
    ")\n",
    "messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"Tell me a story\"),\n",
    "]\n",
    "resp = llm.stream_chat(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4475a6bc-1051-4287-abce-ba83324aeb9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "for r in resp:\n",
    "    print(r.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "009d3f1c-ef35-4126-ae82-0b97adb746e3",
   "metadata": {},
   "source": [
    "## Configure Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e973e3d1-a3c9-43b9-bee1-af3e57946ac3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.bedrock_converse import BedrockConverse\n",
    "\n",
    "llm = BedrockConverse(\n",
    "    model=\"anthropic.claude-3-haiku-20240307-v1:0\",\n",
    "    profile_name=profile_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2c9bcf6-c950-4dfc-abdc-598d5bdedf40",
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = llm.complete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de715fea-4878-4fbb-b415-71250215f3a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bdd4602-e37c-4230-af82-35af5292f9a0",
   "metadata": {},
   "source": [
    "## Connect to Bedrock with Access Keys "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9c80814-6d59-4782-a4bb-cbfcdba6a072",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.bedrock_converse import BedrockConverse\n",
    "\n",
    "llm = BedrockConverse(\n",
    "    model=\"us.amazon.nova-lite-v1:0\",\n",
    "    aws_access_key_id=\"AWS Access Key ID to use\",\n",
    "    aws_secret_access_key=\"AWS Secret Access Key to use\",\n",
    "    aws_session_token=\"AWS Session Token to use\",\n",
    "    region_name=\"AWS Region to use, eg. us-east-1\",\n",
    ")\n",
    "\n",
    "resp = llm.complete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a07982f4-035c-41ca-9dca-49ae6ab3c05a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a73eb402",
   "metadata": {},
   "source": [
    "## Function Calling\n",
    "\n",
    "Claude, Command and Mistral Large models supports native function calling through AWS Bedrock Converse. There's a seamless integration with LlamaIndex tools, through the `predict_and_call` function on the `llm`. \n",
    "\n",
    "This allows the user to attach any tools and let the LLM decide which tools to call (if any).\n",
    "\n",
    "If you wish to perform tool calling as part of an agentic loop, check out our [agent guides](https://docs.llamaindex.ai/en/latest/module_guides/deploying/agents/) instead.\n",
    "\n",
    "**NOTE**: Not all models from AWS Bedrock support function calling and the Converse API. [Check the available features of each LLM here](https://docs.aws.amazon.com/bedrock/latest/userguide/models-features.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "427105db",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.bedrock_converse import BedrockConverse\n",
    "from llama_index.core.tools import FunctionTool\n",
    "\n",
    "\n",
    "def multiply(a: int, b: int) -> int:\n",
    "    \"\"\"Multiple two integers and returns the result integer\"\"\"\n",
    "    return a * b\n",
    "\n",
    "\n",
    "def mystery(a: int, b: int) -> int:\n",
    "    \"\"\"Mystery function on two integers.\"\"\"\n",
    "    return a * b + a + b\n",
    "\n",
    "\n",
    "mystery_tool = FunctionTool.from_defaults(fn=mystery)\n",
    "multiply_tool = FunctionTool.from_defaults(fn=multiply)\n",
    "\n",
    "llm = BedrockConverse(\n",
    "    model=\"anthropic.claude-3-haiku-20240307-v1:0\",\n",
    "    profile_name=profile_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5a2f465",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = llm.predict_and_call(\n",
    "    [mystery_tool, multiply_tool],\n",
    "    user_msg=\"What happens if I run the mystery function on 5 and 7\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41313d40",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37b5d62f",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = llm.predict_and_call(\n",
    "    [mystery_tool, multiply_tool],\n",
    "    user_msg=(\n",
    "        \"\"\"What happens if I run the mystery function on the following pairs of numbers? Generate a separate result for each row:\n",
    "- 1 and 2\n",
    "- 8 and 4\n",
    "- 100 and 20\n",
    "\n",
    "NOTE: you need to run the mystery function for all of the pairs above at the same time \\\n",
    "\"\"\"\n",
    "    ),\n",
    "    allow_parallel_tool_calls=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a8295e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45b6081a",
   "metadata": {},
   "outputs": [],
   "source": [
    "for s in response.sources:\n",
    "    print(f\"Name: {s.tool_name}, Input: {s.raw_input}, Output: {str(s)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7bee3bc",
   "metadata": {},
   "source": [
    "## Async"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f4c144f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.bedrock_converse import BedrockConverse\n",
    "\n",
    "llm = BedrockConverse(\n",
    "    model=\"anthropic.claude-3-haiku-20240307-v1:0\",\n",
    "    aws_access_key_id=\"AWS Access Key ID to use\",\n",
    "    aws_secret_access_key=\"AWS Secret Access Key to use\",\n",
    "    aws_session_token=\"AWS Session Token to use\",\n",
    "    region_name=\"AWS Region to use, eg. us-east-1\",\n",
    ")\n",
    "resp = await llm.acomplete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd72e3a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(resp)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
